# -*- coding: utf-8 -*-
from loguru import logger
from webrunnercore import wr  # wr模块是webrunner内置的脚本增强sdk
from webrunnercore import *


class 测试负载(WebLoadMachine):
    '''
    用户自定义负载信息
    '''
    负载名称 = '默认负载'

    测试机 = [
        {
            "ip地址": '127.0.0.1',
            "端口": 50000,
            "节点数": 1,
            "主节点": True
        },
    ]


class 测试场景(WebScenario):
    '''
    用户自定义场景信息
    '''
    场景名称 = '默认场景'

    模式 = '梯形负载'
    参数 = {
        '用户数': 200,
        '创建速率': 10,
        '运行时长': 60
    }


class Transaction_Morenfenzu(SerialTransaction):
    '''
    事务定义, 一个事务由多个task构成, 每个task只包含一个请求
    '''

    def __init__(self, parent: "User") -> None:
        super().__init__(parent)
        self.news_id = None  # 用于存储新增的新闻ID
        self.token = None    # 用于存储登录token


    @property
    def transaction(self):
        # 事务名称
        return "morenfenzu"


    def on_start(self):
        # 事务启动函数
        super().on_start()


    @task
    def task_0(self):
        url = "http://localhost:8080/api/user/login"
        headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Content-Length': '40', 'Content-Type': 'application/json', 'Host': 'localhost:8080', 'Origin': 'http://localhost:8080', 'Referer': 'http://localhost:8080/login', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"'}
        params = {}
        # application/json (dict)
        data = {}
        data['userName'] = 'admin'
        data['password'] = '123456'

        res = self.post(url, headers=headers, params=params, json=data)
        
        # 检查登录响应
        if res.status_code == 200:
            try:
                response_data = res.json()
                if response_data.get('code') == 200:
                    # 从响应中提取token
                    self.token = response_data.get('data', {}).get('token')
                    if self.token:
                        logger.info(f"登录成功，获取到token: {self.token[:20]}...")
                    else:
                        logger.error("登录响应中未找到token")
                else:
                    logger.error(f"登录失败: {response_data.get('message', '未知错误')}")
            except Exception as e:
                logger.error(f"解析登录响应失败: {e}")
        else:
            logger.error(f"登录请求失败，状态码: {res.status_code}")


    @task
    def task_1(self):
        # 检查是否有token
        if not self.token:
            logger.error("未获取到token，跳过新增新闻操作")
            return
            
        url = "http://localhost:8080/api/news/add"
        headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Content-Type': 'application/json', 'Host': 'localhost:8080', 'Origin': 'http://localhost:8080', 'Referer': 'http://localhost:8080/news', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'token': self.token}
        params = {'currentUserId': '1'}
        
        # 生成唯一的新闻标题，避免重复
        import uuid
        unique_title = f'test_{uuid.uuid4().hex[:8]}'
        data = {
            'newsId': None,
            'newsTitle': unique_title,
            'author': 'test',
            'newsSummary': 'test',
            'newsImagePath': 'tv.jpg',
            'newsContent': f'<p>{unique_title}</p>',
            'creatPersonId': 1
        }

        res = self.post(url, headers=headers, params=params, json=data)
        
        # 检查新增新闻响应
        if res.status_code == 200:
            try:
                response_data = res.json()
                if response_data.get('code') == 200:
                    # 从响应中提取新增的新闻ID
                    if isinstance(response_data.get('data'), dict):
                        self.news_id = response_data.get('data', {}).get('newsId')
                    elif isinstance(response_data.get('data'), str):
                        # 如果data是字符串，需要通过查询获取newsId
                        logger.info("新增新闻成功，但响应格式不包含newsId，将通过查询获取")
                        self.news_id = self._get_news_id_by_title(unique_title)
                    else:
                        self.news_id = response_data.get('newsId')
                    
                    if self.news_id:
                        logger.info(f"新增新闻成功，新闻ID: {self.news_id}")
                    else:
                        logger.warning("新增新闻成功但未获取到新闻ID")
                else:
                    logger.error(f"新增新闻失败: {response_data.get('message', '未知错误')}")
            except Exception as e:
                logger.error(f"解析新增新闻响应失败: {e}")
        else:
            logger.error(f"新增新闻请求失败，状态码: {res.status_code}")

    def _get_news_id_by_title(self, title):
        """通过标题查询获取新闻ID"""
        try:
            url = "http://localhost:8080/api/news/list-page"
            headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Host': 'localhost:8080', 'Referer': 'http://localhost:8080/news', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'token': self.token}
            params = {
                'current': '1',
                'newsTitle': title,  # 使用具体标题查询
                'currentUserId': '1'
            }
            
            res = self.get(url, headers=headers, params=params)
            if res.status_code == 200:
                response_data = res.json()
                if response_data.get('code') == 200:
                    records = response_data.get('data', {}).get('records', [])
                    if records:
                        # 找到匹配标题的新闻
                        for record in records:
                            if record.get('newsTitle') == title:
                                news_id = record.get('newsId')
                                logger.info(f"通过标题'{title}'找到新闻ID: {news_id}")
                                return news_id
                        logger.warning(f"未找到标题为'{title}'的新闻")
                    else:
                        logger.warning("查询结果为空")
                else:
                    logger.error(f"查询新闻失败: {response_data.get('message', '未知错误')}")
            else:
                logger.error(f"查询新闻请求失败，状态码: {res.status_code}")
        except Exception as e:
            logger.error(f"查询新闻异常: {e}")
        return None


    @task
    def task_2(self):
        # 检查是否有token
        if not self.token:
            logger.error("未获取到token，跳过查询操作")
            return
            
        url = "http://localhost:8080/api/news/list-page"
        headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Host': 'localhost:8080', 'Referer': 'http://localhost:8080/news', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'token': self.token}
        params = {
            'current': '1',
            'newsTitle': 'test',  # 查询包含test的新闻
            'currentUserId': '1'
        }
        
        res = self.get(url, headers=headers, params=params)
        
        # 检查查询响应
        if res.status_code == 200:
            try:
                response_data = res.json()
                if response_data.get('code') == 200:
                    logger.info("查询新闻列表成功")
                else:
                    logger.error(f"查询新闻列表失败: {response_data.get('message', '未知错误')}")
            except Exception as e:
                logger.error(f"解析查询响应失败: {e}")
        else:
            logger.error(f"查询新闻列表请求失败，状态码: {res.status_code}")


    @task
    def task_3(self):
        # 检查是否有token和新闻ID
        if not self.token:
            logger.error("未获取到token，跳过更新操作")
            return
            
        if not self.news_id:
            logger.error("未获取到新闻ID，跳过更新操作")
            return
            
        url = "http://localhost:8080/api/news/update/true"
        headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Content-Type': 'application/json', 'Host': 'localhost:8080', 'Origin': 'http://localhost:8080', 'Referer': 'http://localhost:8080/news', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'token': self.token}
        data = {
            'newsId': self.news_id,  # 使用新增的新闻ID
            'newsTitle': f'test1_{self.news_id}',  # 标题包含ID便于识别
            'author': 'test',
            'newsSummary': 'test',
            'newsImagePath': 'tv.jpg',
            'newsContent': f'<p>test1_{self.news_id}</p>',
            'creatPersonId': 1
        }

        res = self.put(url, headers=headers, json=data)
        
        # 检查更新响应
        if res.status_code == 200:
            try:
                response_data = res.json()
                if response_data.get('code') == 200:
                    logger.info(f"更新新闻成功，新闻ID: {self.news_id}")
                else:
                    logger.error(f"更新新闻失败: {response_data.get('message', '未知错误')}")
            except Exception as e:
                logger.error(f"解析更新响应失败: {e}")
        else:
            logger.error(f"更新新闻请求失败，状态码: {res.status_code}")


    @task
    def task_4(self):
        # 检查是否有token和新闻ID
        if not self.token:
            logger.error("未获取到token，跳过删除操作")
            return
            
        if not self.news_id:
            logger.error("未获取到新闻ID，跳过删除操作")
            return
            
        url = "http://localhost:8080/api/news/delete"
        headers = {'Accept': 'application/json, text/plain, */*', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Connection': 'keep-alive', 'Host': 'localhost:8080', 'Origin': 'http://localhost:8080', 'Referer': 'http://localhost:8080/news', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-origin', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36', 'sec-ch-ua': '"Not.A/Brand";v="8", "Chromium";v="114"', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'token': self.token}
        params = {'newsId': self.news_id}  # 使用新增的新闻ID

        res = self.delete(url, headers=headers, params=params)
        
        # 检查删除响应
        if res.status_code == 200:
            try:
                response_data = res.json()
                if response_data.get('code') == 200:
                    logger.info(f"删除新闻成功，新闻ID: {self.news_id}")
                else:
                    logger.error(f"删除新闻失败: {response_data.get('message', '未知错误')}")
            except Exception as e:
                logger.error(f"解析删除响应失败: {e}")
        else:
            logger.error(f"删除新闻请求失败，状态码: {res.status_code}")


    def on_stop(self):
        # 事务结束函数
        super().on_stop()


class WebrunnerAction(BrowserAction):
    '''
    事务集合, 一个Action包含多个Transaction事务, @task装饰器参数表示事务混合比
    '''
    def __init__(self, parent: "User") -> None:
        super().__init__(parent)

    def on_start(self):
        super().on_start()

    @task(1)
    @transaction("morenfenzu")
    def task_morenfenzu(self):
        # 执行事务
        Transaction_Morenfenzu(self).run()

    def on_stop(self):
        super().on_stop()


class WebrunnerUser(CFastHttpUser):
    '''
    虚拟用户, 一个用户循环执行一个Action
    '''
    host = ""
    tasks = [WebrunnerAction]

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    def on_start(self):
        # 所有虚拟用户创建完成后开始执行，主要用于定义参数化和检查点的策略
        super().on_start()